/* diffJac2d.c                               (\(\
 * Jannis Teunnissen and Florian Speelman    (='.')
 * Scientific Computing - spring 2010        o(_")")
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>

#include "getDiffPars.h"
#include "divideElems.h"

#define WRITE_DATFILE 1
#define WRITE_PLOTS 2

double diffSolver(diffPars *p, MPI_Comm *commPtr);
void stepDiff(double **c, diffPars *p, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *commPtr);
void initConcentration(double **c, int nPoints, int totalPoints, int myrank, int ntasks);
void initPars(int argc, char *argv[], diffPars *p, MPI_Comm *commPtr);
void writeDiff(double **c, double time, diffPars *p, int mode, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *commPtr);

int main(int argc, char *argv[]) {
    int rc;
    
    rc = MPI_Init(&argc, &argv);
    if(rc != MPI_SUCCESS) {
        printf("MPI_Init error %d\n", rc);
        MPI_Abort(MPI_COMM_WORLD, rc);
    }
    
    MPI_Comm comm;
    diffPars p;
    
    initPars(argc, argv, &p, &comm);
    
    double computeTime;
    computeTime = diffSolver(&p, &comm);
    
    if(computeTime > 0.) printf("Total compute time %e seconds\n", computeTime);
    
    MPI_Finalize();
    return 0;
}

double diffSolver(diffPars *p, MPI_Comm *commPtr) {
    int myrank, ntasks;
    
    MPI_Comm_size(*commPtr, &ntasks);
    MPI_Comm_rank(*commPtr, &myrank);
    
    int nElemsArray[ntasks];    //Array containing the number of elements of each task,
    int displs[ntasks];         //Displacement array
    
    {
        int myWeight = 1;
        if(p->loadBalance == 1) {
            myWeight = benchmark();
        }
        int weights[ntasks];
        MPI_Gather(&myWeight, 1, MPI_INT, weights, 1, MPI_INT, 0, *commPtr);
        if(myrank == 0) {
            printWeights(weights, ntasks);
            divideElems(nElemsArray, displs, weights, p->totalPoints, ntasks);
        }
    }
    
    //Now share nElemsArray and displs to later on use allgatherv
    MPI_Bcast(nElemsArray, ntasks, MPI_INT, 0, *commPtr);
    MPI_Bcast(displs, ntasks, MPI_INT, 0, *commPtr);
    
    int nPoints = nElemsArray[myrank];
    double **c = malloc((nPoints + 2) * sizeof(double*));
    for(int i = 0; i < nPoints + 2; i++) {
        c[i] = malloc((p->totalPoints + 2) * sizeof(double));
    }

    initConcentration(c, nPoints, p->totalPoints, myrank, ntasks);

    double startTime = 0., endTime = 0.;
    
    MPI_Barrier(*commPtr);
    
    if(myrank == 0) startTime = MPI_Wtime();
    stepDiff(c, p, nElemsArray, displs, myrank, ntasks, commPtr);
    if(myrank == 0) endTime = MPI_Wtime();
    
    return endTime - startTime;
}

static inline void updatePoint(double **cNew, double **c, double *maxDiff, int i, int j, int checkDiff) {
	double newval;
	double difference;
	newval = 0.25 * (c[i-1][j] + c[i+1][j] + c[i][j-1] + c[i][j+1]);
	if(checkDiff) {
		difference = fabs(newval - c[i][j]);
		if(difference > *maxDiff) *maxDiff = difference;
	}
	cNew[i][j] = newval;
}

void stepDiff(double **c, diffPars *p, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *commPtr) {
    int i, j;
    int up, down;
    MPI_Cart_shift(*commPtr, 0, 1, &up, &down);
    
    int nTimeSteps = p->maxTime / p->dt;
    int timePerOutput;
    int outputCounter = 0;
    
    if(p->nOutputs > 0) {
        timePerOutput = nTimeSteps / p->nOutputs;
    } else {
        timePerOutput = 1e7;     //No output
        outputCounter = 1;
    }
    
    int nPoints = nElemsArray[myrank];
	int endPoint = nPoints + 1;
	int startPoint = 1;
    int N = p->totalPoints;
	int its = 0;
	int checkInterval = p->checkInterval;
    double dFactor = p->dt * p->D / (p->dx * p->dx);
    double dFactor2 = 1. - 4. * dFactor;
    double maxDiff;
	
    double **cNew = malloc((nPoints + 2) * sizeof(double*));
    for(i = 0; i < nPoints + 2; i++) {
        cNew[i] = malloc((N + 2) * sizeof(double));
    }
	
	if(myrank == ntasks - 1) {
		endPoint--;
		for(j = 0; j < N + 2; j++) {
			cNew[endPoint][j] = c[endPoint][j];
		}
	}
	
	if(myrank == 0) {
		startPoint++;
		for(j = 0; j < N + 2; j++) {
			cNew[1][j] = c[1][j];
		}
	}
    
    MPI_Request rq[4];

    while(1) {
		maxDiff = 0.0;
		its++;
        int checkDiff = (its % checkInterval == 0);
		if(its >= outputCounter * timePerOutput) {
			writeDiff(c, its, p, WRITE_DATFILE, nElemsArray, displs, myrank, ntasks, commPtr);
			outputCounter++;
		}
		
		for(i = startPoint; i < endPoint; i++) {
            c[i][N+1] = c[i][1];
            c[i][0] = c[i][N];
        }
		
		for(j = 1; j < N + 1; j++) {
			updatePoint(cNew, c, &maxDiff, startPoint, j, checkDiff);
		}
		
		for(j = 1; j < N + 1; j++) {
			updatePoint(cNew, c, &maxDiff, endPoint-1, j, checkDiff);
		}
		
		MPI_Irecv(cNew[nPoints+1], N+2, MPI_DOUBLE, down, 0, *commPtr, &rq[0]);
		MPI_Irecv(cNew[0], N+2, MPI_DOUBLE, up, 1, *commPtr, &rq[1]);
		MPI_Isend(cNew[1], N+2, MPI_DOUBLE, up, 0, *commPtr, &rq[2]);
		MPI_Isend(cNew[nPoints], N+2, MPI_DOUBLE, down, 1, *commPtr, &rq[3]);
		
        
		for(i = startPoint + 1; i < endPoint - 1; i++) {
            for(j = 1; j < N + 1; j++) {
                updatePoint(cNew, c, &maxDiff, i, j, checkDiff);
            }
        }
		
        MPI_Waitall(4, rq, MPI_STATUS_IGNORE);
        if(checkDiff) {
			double globalMaxDiff;
			MPI_Allreduce(&maxDiff, &globalMaxDiff, 1, MPI_DOUBLE, MPI_MAX, *commPtr);
			if(globalMaxDiff < p->threshold) {
				break;
			}
		}
		
        {
            double **tempPtr = c;
            c = cNew;
            cNew = tempPtr;
        }
    }
	
	if(myrank == 0)
		printf("Converged after %i iterations\n", its);

    if(p->imgOutput == 1) {
        writeDiff(c, 0, p, WRITE_PLOTS, nElemsArray, displs, myrank, ntasks, commPtr);
    }
    free(c);
    free(cNew);
}

void writeDiff(double **c, double time, diffPars *p, int mode, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *commPtr) {
    static int counter = 0;
    static double **cComplete = NULL;
    int N = p->totalPoints;

    if(counter == 0 && myrank == 0) {
        FILE *datDir = fopen("dat/.dircheck", "w");
        FILE *pngDir = fopen("png/.dircheck", "w");
        FILE *outputInfo = fopen("dat/outputInfo.txt", "w");
        if(pngDir == NULL || datDir == NULL) {
            if(myrank == 0) fprintf(stderr, "Make sure directories 'dat' and 'png' exist\n");
            MPI_Abort(*commPtr, 0);
        }
        fclose(datDir); fclose(pngDir); fclose(outputInfo);
        cComplete = malloc(N * sizeof(double*));        //Don't free cComplete between calls for performance
        for(int i = 0; i < N; i++) {
            cComplete[i] = malloc(N * sizeof(double));
        }
    }
    
    if(mode == WRITE_DATFILE) {
        if(myrank == 0) {
            for(int row = 0; row < nElemsArray[0]; row++) {
                for(int j = 0; j < N; j++) {
                    cComplete[row][j] = c[row+1][j+1];
                }
            }
            for(int rank = 1; rank < ntasks; rank++) {
                for(int row = 0; row < nElemsArray[rank]; row++) {
                    MPI_Recv(cComplete[displs[rank] + row], N, MPI_DOUBLE, rank, rank, *commPtr, MPI_STATUS_IGNORE);
                }
            }
        } else {
            for(int row = 1; row < nElemsArray[myrank]+1; row++) {
                MPI_Send(&(c[row][1]), N, MPI_DOUBLE, 0, myrank, *commPtr);
            }
        }
        
        if(myrank == 0) {
            char filename[40];
            sprintf(filename, "dat/diff_%d.dat", counter);
            FILE *output = fopen(filename, "w");
            FILE *outputInfo = fopen("dat/outputInfo.txt", "a");
            fprintf(outputInfo, "%e\t%e\n", time, 1337.);
            for(int i = 0; i < N; i++) {
                for(int j = 0; j < N; j++) {
                    fprintf(output, "%e\t%e\t%e\n", p->dx * i, p->dx * j, cComplete[i][j]);
                }
                fprintf(output, "\n");
            }
            fclose(output); fclose(outputInfo);
        }
        counter++;
    }
    
    if (mode == WRITE_PLOTS && myrank == 0) {
        FILE *gnuplot = popen("gnuplot -persist", "w");
        FILE *outputInfo = fopen("dat/outputInfo.txt", "r");
        fprintf(gnuplot,"set nokey\n");
        fprintf(gnuplot,"set terminal png medium\n");
        
        fprintf(gnuplot, "set title 'Absolute error over time'\n");
        fprintf(gnuplot, "set output 'png/diffError.png'\n");
        fprintf(gnuplot, "plot 'dat/outputInfo.txt' using 1:2 w l\n");
        fprintf(gnuplot, "set pm3d\n");
        fprintf(gnuplot, "set hidden3d\n");
        fprintf(gnuplot,"set xlabel 'x'\n");
        fprintf(gnuplot,"set ylabel 'y'\n");
        fprintf(gnuplot,"set zlabel 'c'\n");

        double error;
        for(int i = 0; i < counter; i++) {
            if(fscanf(outputInfo, "%le %le\n", &time, &error)) {
                fprintf(gnuplot, "set title 'Concentration at t = %.2lf s, max error %.2e'\n", time, error);
                fprintf(gnuplot, "set output 'png/diff_%06d.png'\n", i);
                fprintf(gnuplot, "splot 'dat/diff_%d.dat' using 2:1:3 w l\n", i);
            }
        }
        fclose(gnuplot); fclose(outputInfo);
    }
}

void initConcentration(double **c, int nPoints, int totalPoints, int myrank, int ntasks) {
    int i, j;
    
    for(i = 0; i < nPoints + 2; i++) {
        for(j = 0; j < totalPoints + 2; j++) {
            c[i][j] = 0.;
        }
    }
    if(myrank == ntasks-1) {
        for(j = 0; j < totalPoints + 2; j++) {
            c[nPoints][j] = 1.;
        }
    }
}


